#!/usr/bin/env python
#author Matthew M Hindle
#email matthew.hindle@ed.ac.uk
#version 0.1
#python_version Python 2.7.3
#usage python CoPSAPfamSearch.py dirwifasta --cpus 1 --outDir myDir --cuttoff cut_ga
#licence GPLv3
import os, sys
import subprocess, multiprocessing
import select
import argparse

#the number of cpus forms the default for the number of threads
cpus = multiprocessing.cpu_count()

parser = argparse.ArgumentParser(description='Runs parallel hmmsearch jobs on a fasta file.')
parser.add_argument('inDir', type=str, help='Dir of FASTA files')
parser.add_argument('--outDir', dest='outDir', type=str, help='Dir to put results', default=os.path.abspath('PFAMOUT'))
parser.add_argument('---hmmFile', dest='hmmFile', type=str, help='The path to a hmm e.g. Pfam-A.hmm from ftp://ftp.sanger.ac.uk/pub/databases/Pfam/current_release. We reccomend you index the file first with hmmpress', default='/home/mhindle/Pfam07_01_2013/Pfam-A.hmm')
parser.add_argument('--cpus', dest='cpus', type=int, help='CPUS to use', default=cpus)
parser.add_argument('--cuttoff', dest='cuttoff', type=str, help='Threshold to use', default='cut_ga')
parser.add_argument('--max', dest='max', action="store_true", help='Turn of huristics', default=False)
parser.add_argument('--notnice', dest='nice', action="store_false", help='Do not nice processes on linux', default=False)

args = parser.parse_args()

outDir = os.path.abspath(args.outDir)
if not os.path.exists(outDir):
    os.makedirs(outDir)

poller = select.epoll() #our external process queue
subprocs = {} #jobs to do
finSubprocs = {} #jobs we have done

#start hmmsearch jobs on fasta files
for fileIn in os.listdir(args.inDir):
	if fileIn.endswith(".fasta"):
		fastaFile = os.path.abspath(args.inDir+'/'+fileIn)

		print 'Running hmmsearch on '+fileIn
		outputName = args.outDir+'/'+os.path.splitext(fileIn)[0]

		fileOut = outputName+'HMM.out'
		tableSeqOut = outputName+'SeqHMM.out'
		tableDomOut = outputName+'DomHMM.out'
		alignOut  = outputName+'AlignHMM.out'

		command = [
			'hmmsearch',
			'--cpu', str(args.cpus),
			'--'+args.cuttoff, 
			'-o', fileOut,
			'-A', alignOut,
			'--tblout', tableSeqOut,
			'--domtblout', tableDomOut,
			args.hmmFile, fastaFile]
		
		#turn of huristics?
		if args.max:
			command.append('--max')

		#I assume people will want to nice hmmr processes on linux but if they don'e we support that?
		if args.nice and sys.platform.startswith('linux'):
			command.insert(0, 'nice')

		subproc = subprocess.Popen(command,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
		subprocs[subproc.stdout.fileno()] = subproc
		poller.register(subproc.stdout, select.EPOLLHUP )

#wait for results and output any stdout and stderr
while len(subprocs) > len(finSubprocs):
	#print str(len(finSubprocs))+' of '+str(len(subprocs))
	for fd, flags in poller.poll(timeout=1): #1 second wait interval for jobs to finish
		done_proc = subprocs[fd]
		poller.unregister(fd)
		finSubprocs[fd] = done_proc
		for line in done_proc.stdout:
			print 'STDOUT: '+line
		for line in done_proc.stderr:
			print 'STDERR: '+line
